#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Rust知识库扩充工具
自动从公开资源和模板生成Rust知识库条目
"""

import json
import os
import random
import time
# import requests  # 可选，目前未使用
from typing import List, Dict
import re

class RustKnowledgeBaseExpander:
    def __init__(self, output_path='rust_knowledge_base/rust_docs_expanded.json'):
        self.output_path = output_path
        self.existing_knowledge = []
        self.new_entries = []
        self.categories = ["definition", "usage", "error_debug", "comparison", "faq"]
        self.difficulties = ["beginner", "intermediate", "advanced"]
        
        # Rust主题分类和相关子主题
        self.rust_topics = {
            "所有权与生命周期": ["所有权系统", "借用规则", "生命周期参数", "移动语义", "复制语义", "引用计数"],
            "基础类型": ["标量类型", "复合类型", "元组", "数组", "切片", "字符串"],
            "函数与闭包": ["函数定义", "函数参数", "返回值", "闭包", "高阶函数", "递归"],
            "结构体与枚举": ["结构体定义", "方法实现", "枚举类型", "模式匹配", "Option类型", "Result类型"],
            "特质系统": ["特质定义", "默认实现", "特质边界", "关联类型", "特质对象", "孤儿规则"],
            "泛型编程": ["泛型函数", "泛型类型", "泛型特质", "泛型约束", "类型参数", "单态化"],
            "模块系统": ["模块定义", "可见性", "use声明", "pub关键字", "super和self", "包管理"],
            "错误处理": ["Result类型", "?操作符", "panic宏", "错误传播", "自定义错误", "错误转换"],
            "集合类型": ["Vec动态数组", "HashMap哈希表", "HashSet集合", "BTreeMap有序映射", "BTreeSet有序集合", "迭代器"],
            "智能指针": ["Box智能指针", "Rc引用计数", "Arc原子引用", "RefCell内部可变性", "Cell类型", "Weak引用"],
            "并发编程": ["线程创建", "消息传递", "互斥锁", "读写锁", "原子类型", "线程安全"],
            "异步编程": ["async/await", "Future特质", "异步运行时", "异步任务", "并发控制", "流处理"],
            "不安全Rust": ["unsafe关键字", "原始指针", "FFI接口", "内联汇编", "全局变量", "内存布局"],
            "宏系统": ["声明宏", "过程宏", "derive宏", "属性宏", "函数式宏", "宏展开"],
            "测试与基准": ["单元测试", "集成测试", "文档测试", "基准测试", "测试属性", "测试工具"],
            "性能优化": ["零成本抽象", "内存优化", "算法优化", "并发优化", "编译优化", "分析工具"],
            "Rust生态": ["Cargo工具", "crates.io", "标准库", "第三方库", "开发工具", "社区实践"],
            "高级特性": ["无栈协程", "类型族", "常量泛型", "稳定ABI", "内联汇编", "链接优化"]
        }
        
        # 代码模板
        self.code_templates = {
            "definition": [
                "// 定义{topic}\n{code}",
                "// {topic}的基本实现\n{code}",
                "// {topic}的核心概念示例\n{code}"
            ],
            "usage": [
                "// 使用{topic}\n{code}",
                "// {topic}的常见用法\n{code}",
                "// {topic}的实际应用示例\n{code}"
            ],
            "error_debug": [
                "// {topic}的错误示例\n{code}\n// 修复方法: {fix}",
                "// 错误代码\n{code}\n// 正确版本\n{correct}",
                "// 编译错误: {error}\n{code}"
            ],
            "comparison": [
                "// Rust中的{topic}\n{rust_code}\n\n// 其他语言中的类似实现\n// {other_language}:\n// {other_code}",
                "// {topic}在Rust vs {other_language}\n\n// Rust:\n{rust_code}\n\n// {other_language}:\n// {other_code}"
            ],
            "faq": [
                "// 如何在Rust中实现{topic}？\n{code}",
                "// {topic}的最佳实践\n{code}",
                "// 处理{topic}的推荐方法\n{code}"
            ]
        }
    
    def load_existing_knowledge(self, existing_path='rust_knowledge_base/rust_docs_sample.json'):
        """加载现有的知识库文件"""
        if os.path.exists(existing_path):
            try:
                with open(existing_path, 'r', encoding='utf-8') as f:
                    self.existing_knowledge = json.load(f)
                print(f"成功加载现有知识库，共{len(self.existing_knowledge)}条记录")
                return True
            except Exception as e:
                print(f"加载现有知识库失败: {e}")
        return False
    
    def generate_entry_from_template(self, entry_id: str, category: str, difficulty: str) -> Dict:
        """从模板生成知识条目"""
        # 随机选择主题和子主题
        main_topic = random.choice(list(self.rust_topics.keys()))
        sub_topic = random.choice(self.rust_topics[main_topic])
        
        # 根据类别和难度生成内容
        if category == "definition":
            content = f"Rust中的{sub_topic}是{main_topic}的重要组成部分。{sub_topic}用于{self._generate_description(category, sub_topic)}。"
            code = self._generate_code_snippet(category, sub_topic, difficulty)
        elif category == "usage":
            content = f"在Rust中使用{sub_topic}的方法：{self._generate_description(category, sub_topic)}。{self._generate_usage_examples(sub_topic, difficulty)}。"
            code = self._generate_code_snippet(category, sub_topic, difficulty)
        elif category == "error_debug":
            content = f"常见的{sub_topic}错误及解决方案：{self._generate_description(category, sub_topic)}。{self._generate_error_fix(sub_topic)}。"
            code = self._generate_code_snippet(category, sub_topic, difficulty)
        elif category == "comparison":
            other_language = random.choice(["C++", "Python", "Go", "JavaScript"])
            content = f"Rust与{other_language}中的{sub_topic}对比：{self._generate_description(category, sub_topic)}。{self._generate_language_comparison(sub_topic, other_language)}。"
            code = self._generate_code_snippet(category, sub_topic, difficulty, other_language)
        elif category == "faq":
            content = f"关于Rust中{sub_topic}的常见问题：{self._generate_description(category, sub_topic)}。{self._generate_faq_answer(sub_topic)}。"
            code = self._generate_code_snippet(category, sub_topic, difficulty)
        
        # 生成标签
        tags = [main_topic, sub_topic]
        # 根据内容添加一些相关标签
        if "所有权" in content:
            tags.append("所有权")
        if "生命周期" in content:
            tags.append("生命周期")
        if "借用" in content:
            tags.append("借用")
        if "并发" in content:
            tags.append("并发")
        if "异步" in content:
            tags.append("异步")
        
        # 确保标签不重复且数量适中
        tags = list(set(tags))[:5]  # 最多5个标签
        
        # 选择合适的代码模板
        code_template = random.choice(self.code_templates[category])
        
        # 创建参数字典，始终包含topic
        params = {'topic': sub_topic}
        
        # 根据模板字符串中需要的参数，添加相应的值
        if '{code}' in code_template:
            params['code'] = code
        
        if '{error}' in code_template:
            params['error'] = code.get('error', code) if isinstance(code, dict) else code
        
        if '{fix}' in code_template:
            params['fix'] = code.get('fix', '修复相关错误') if isinstance(code, dict) else '修复相关错误'
        
        if '{correct}' in code_template:
            params['correct'] = code.get('correct', '// 正确代码示例') if isinstance(code, dict) else '// 正确代码示例'
        
        if '{rust_code}' in code_template:
            params['rust_code'] = code.get('rust', code) if isinstance(code, dict) else code
        
        if '{other_language}' in code_template:
            params['other_language'] = code.get('other_language', 'C++') if isinstance(code, dict) else 'C++'
        
        if '{other_code}' in code_template:
            params['other_code'] = code.get('other_code', '// 其他语言代码示例') if isinstance(code, dict) else '// 其他语言代码示例'
        
        # 格式化模板
        filled_code = code_template.format(**params)
        
        return {
            "id": entry_id,
            "title": sub_topic,
            "content": content,
            "code": filled_code,
            "tags": tags,
            "category": category,
            "difficulty": difficulty
        }
    
    def _generate_description(self, category: str, topic: str) -> str:
        """生成描述文本"""
        descriptions = {
            "definition": f"{topic}是Rust语言的基础概念，它提供了{self._get_topic_feature(topic)}的能力",
            "usage": f"正确使用{topic}可以提高代码的{self._get_benefit(topic)}",
            "error_debug": f"常见的错误包括{self._get_common_errors(topic)}",
            "comparison": f"Rust的{topic}实现更加安全高效",
            "faq": f"这个问题经常被Rust初学者问到"
        }
        return descriptions.get(category, "这是一个重要的Rust概念")
    
    def _generate_code_snippet(self, category: str, topic: str, difficulty: str, other_language: str = "C++") -> str:
        """生成代码片段"""
        # 根据主题和类别生成相应的代码片段
        if topic == "所有权系统":
            return "let s1 = String::from(\"hello\");\nlet s2 = s1; // s1的所有权移动到s2\n// println!(\"{}\", s1); // 错误：s1不再有效"
        elif topic == "借用规则":
            return "let mut s = String::from(\"hello\");\nlet r1 = &s; // 不可变借用\nlet r2 = &s; // 不可变借用\n// let r3 = &mut s; // 错误：不能同时有可变和不可变借用"
        elif topic == "生命周期参数":
            return "fn longest<'a>(x: &'a str, y: &'a str) -> &'a str {\n    if x.len() > y.len() { x } else { y }\n}"
        elif topic == "结构体定义":
            return "struct User {\n    name: String,\n    age: u32,\n    active: bool,\n}\n\nlet user = User {\n    name: String::from(\"Alice\"),\n    age: 30,\n    active: true,\n};"
        elif topic == "错误处理" and category == "error_debug":
            return {
                "error": "fn read_file() {\n    let file = File::open(\"hello.txt\"); // 缺少错误处理\n    let mut contents = String::new();\n    file.read_to_string(&mut contents); // 错误\n}",
                "correct": "use std::fs::File;\nuse std::io::{self, Read};\n\nfn read_file() -> io::Result<String> {\n    let mut file = File::open(\"hello.txt\")?;\n    let mut contents = String::new();\n    file.read_to_string(&mut contents)?;\n    Ok(contents)\n}",
                "fix": "使用?操作符进行错误传播"
            }
        elif topic == "并发编程" and category == "comparison":
            return {
                "rust": "use std::thread;\nuse std::sync::mpsc;\n\nfn main() {\n    let (tx, rx) = mpsc::channel();\n    \n    thread::spawn(move || {\n        tx.send(String::from(\"hello\")).unwrap();\n    });\n    \n    let received = rx.recv().unwrap();\n    println!(\"Got: {}\", received);\n}",
                "other_language": other_language,
                "other_code": "// C++版本\n#include <thread>\n#include <iostream>\n#include <string>\n\nint main() {\n    std::thread t([]{\n        std::cout << \"hello\" << std::endl;\n    });\n    t.join();\n    return 0;\n}"
            }
        elif topic == "闭包":
            return "// 闭包定义\nlet add_one = |x: i32| -> i32 { x + 1 };\nlet result = add_one(5); // 结果是6\n\n// 更简洁的语法\nlet multiply = |x, y| x * y;\nlet product = multiply(3, 4); // 结果是12"
        elif topic == "泛型函数":
            return "// 泛型函数定义\nf<T>(x: T, y: T) -> T where T: std::ops::Add<Output = T> {\n    x + y\n}\n\n// 使用泛型函数\nlet sum_i32 = add(5, 10); // i32类型\nlet sum_f64 = add(3.14, 2.71); // f64类型"
        elif topic == "特质定义":
            return "// 特质定义\ntrait Animal {\n    fn name(&self) -> &str;\n    \n    // 默认实现\n    fn speak(&self) -> String {\n        format!(\"{} makes a sound\", self.name())\n    }\n}\n\n// 实现特质\nstruct Dog { name: String }\nimpl Animal for Dog {\n    fn name(&self) -> &str { &self.name }\n    fn speak(&self) -> String {\n        format!(\"{} barks\", self.name())\n    }\n}"
        else:
            # 为其他主题生成通用代码
            return self._generate_generic_code(topic, difficulty)
    
    def _generate_generic_code(self, topic: str, difficulty: str) -> str:
        """生成通用代码片段"""
        if difficulty == "beginner":
            return f"// {topic}的基础示例\nfn main() {{\n    println!(\"Hello, {topic}!\");\n}}"
        elif difficulty == "intermediate":
            return f"// {topic}的中级示例\nfn process_{topic.lower().replace(' ', '_')}() -> Result<(), Box<dyn std::error::Error>> {{\n    // 实现{topic}的逻辑\n    Ok(())\n}}\n\nfn main() {{\n    if let Err(e) = process_{topic.lower().replace(' ', '_')}() {{\n        eprintln!(\"Error: {{}}\", e);\n    }}\n}}"
        else: # advanced
            # 先选择特性
            sync_feature = random.choice(['Mutex', 'Arc', 'RwLock'])
            data_feature = random.choice(['Mutex', 'Arc'])
            # 然后构建字符串
            code = f"// {topic}的高级示例\nuse std::sync::{sync_feature};\nuse std::thread;\n\n"
            code += f"fn advanced_{topic.lower().replace(' ', '_')}() {{\n    // 复杂的{topic}实现\n    let data = {data_feature}::new(vec![1, 2, 3]);\n    // 多线程或其他高级特性\n}}\n\n"
            code += f"fn main() {{\n    advanced_{topic.lower().replace(' ', '_')}();\n}}"
            return code
    
    def _get_topic_feature(self, topic: str) -> str:
        """获取主题特性描述"""
        features = [
            "内存安全", "并发控制", "类型安全", "零成本抽象", "高性能",
            "错误处理", "代码复用", "模块化", "可扩展性", "可读性"
        ]
        return random.choice(features)
    
    def _get_benefit(self, topic: str) -> str:
        """获取使用主题的好处"""
        benefits = [
            "安全性", "性能", "可维护性", "可读性", "可扩展性",
            "代码复用", "错误预防", "运行效率", "开发效率", "可靠性"
        ]
        return random.choice(benefits)
    
    def _get_common_errors(self, topic: str) -> str:
        """获取常见错误描述"""
        errors = [
            "生命周期不匹配", "借用规则违反", "类型错误", "空指针解引用",
            "数据竞争", "资源泄漏", "无限递归", "溢出", "死锁", "错误处理不当"
        ]
        return random.choice(errors)
    
    def _generate_usage_examples(self, topic: str, difficulty: str) -> str:
        """生成使用示例描述"""
        examples = [
            f"在实际项目中，{topic}常用于处理{self._get_use_case(topic)}",
            f"{topic}的典型应用场景包括{self._get_use_case(topic)}",
            f"推荐在{self._get_situation(topic)}时使用{topic}"
        ]
        return random.choice(examples)
    
    def _generate_error_fix(self, topic: str) -> str:
        """生成错误修复方法"""
        fixes = [
            "检查生命周期注解", "遵循借用规则", "正确使用智能指针",
            "添加适当的类型约束", "实现错误处理逻辑", "优化内存使用"
        ]
        return random.choice(fixes)
    
    def _generate_language_comparison(self, topic: str, other_language: str) -> str:
        """生成语言对比描述"""
        comparisons = [
            f"Rust的{topic}比{other_language}更安全",
            f"Rust在{topic}方面提供了编译时检查，而{other_language}依赖运行时",
            f"Rust的{topic}实现兼顾了安全和性能，{other_language}在这方面有所取舍"
        ]
        return random.choice(comparisons)
    
    def _generate_faq_answer(self, topic: str) -> str:
        """生成FAQ答案"""
        answers = [
            f"要实现{topic}，建议遵循Rust的最佳实践",
            f"{topic}的关键在于理解Rust的所有权模型",
            f"对于{topic}，推荐使用标准库中的{self._get_recommended_feature(topic)}"
        ]
        return random.choice(answers)
    
    def _get_use_case(self, topic: str) -> str:
        """获取使用场景"""
        use_cases = [
            "复杂数据结构", "并发程序", "系统编程", "Web后端",
            "命令行工具", "游戏开发", "嵌入式系统", "高性能计算"
        ]
        return random.choice(use_cases)
    
    def _get_situation(self, topic: str) -> str:
        """获取适用情况"""
        situations = [
            "需要高性能", "对安全性要求高", "处理复杂数据",
            "并发编程", "系统级开发", "跨平台应用"
        ]
        return random.choice(situations)
    
    def _get_recommended_feature(self, topic: str) -> str:
        """获取推荐特性"""
        features = [
            "标准集合类型", "智能指针", "并发原语", "错误处理机制",
            "特质系统", "泛型", "异步运行时", "宏系统"
        ]
        return random.choice(features)
    
    def fetch_from_rust_docs(self, topic: str) -> Dict:
        """尝试从Rust官方文档获取内容（可选增强功能）"""
        # 这是一个占位函数，可以扩展为实际的API调用
        # 由于Rust官方文档可能没有开放API，可以考虑使用其他公开资源
        return {}
    
    def generate_knowledge_base(self, target_count: int = 500):
        """生成指定数量的知识条目"""
        # 确定需要生成的新条目数量
        existing_count = len(self.existing_knowledge)
        if existing_count >= target_count:
            print(f"现有知识库已有{existing_count}条记录，已满足目标数量{target_count}")
            return
        
        new_count = target_count - existing_count
        print(f"需要生成{new_count}条新的知识条目")
        
        # 生成新条目
        for i in range(new_count):
            # 生成唯一ID
            entry_id = f"k{i + existing_count + 1:03d}"
            
            # 随机选择类别和难度，但保持一定的分布
            category_weights = {
                "definition": 0.3,
                "usage": 0.3,
                "error_debug": 0.2,
                "comparison": 0.1,
                "faq": 0.1
            }
            category = random.choices(
                list(category_weights.keys()),
                weights=list(category_weights.values())
            )[0]
            
            # 根据类别调整难度分布
            if category in ["definition", "usage"]:
                difficulty_weights = {"beginner": 0.4, "intermediate": 0.5, "advanced": 0.1}
            else:
                difficulty_weights = {"beginner": 0.2, "intermediate": 0.5, "advanced": 0.3}
            
            difficulty = random.choices(
                list(difficulty_weights.keys()),
                weights=list(difficulty_weights.values())
            )[0]
            
            # 生成条目
            entry = self.generate_entry_from_template(entry_id, category, difficulty)
            self.new_entries.append(entry)
            
            # 显示进度
            if (i + 1) % 50 == 0:
                print(f"已生成{i + 1}/{new_count}条记录")
            
            # 避免生成过快
            time.sleep(0.01)
        
        print(f"成功生成{len(self.new_entries)}条新的知识条目")
    
    def combine_and_save(self):
        """合并现有条目和新条目并保存"""
        # 合并所有条目
        combined_knowledge = self.existing_knowledge + self.new_entries
        
        # 确保ID唯一
        seen_ids = set()
        unique_knowledge = []
        for entry in combined_knowledge:
            if entry["id"] not in seen_ids:
                seen_ids.add(entry["id"])
                unique_knowledge.append(entry)
        
        # 按ID排序
        unique_knowledge.sort(key=lambda x: x["id"])
        
        # 保存到文件
        os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
        with open(self.output_path, 'w', encoding='utf-8') as f:
            json.dump(unique_knowledge, f, ensure_ascii=False, indent=2)
        
        print(f"成功保存扩充后的知识库到{self.output_path}")
        print(f"最终知识库包含{len(unique_knowledge)}条记录")
        
        # 统计各类别的分布
        category_count = {}
        difficulty_count = {}
        for entry in unique_knowledge:
            category = entry.get("category", "unknown")
            difficulty = entry.get("difficulty", "unknown")
            category_count[category] = category_count.get(category, 0) + 1
            difficulty_count[difficulty] = difficulty_count.get(difficulty, 0) + 1
        
        print("\n类别分布:")
        for cat, count in category_count.items():
            print(f"  {cat}: {count}")
        
        print("\n难度分布:")
        for diff, count in difficulty_count.items():
            print(f"  {diff}: {count}")
    
    def validate_knowledge_base(self, file_path=None):
        """验证知识库格式是否正确"""
        path_to_check = file_path or self.output_path
        if not os.path.exists(path_to_check):
            print(f"文件不存在: {path_to_check}")
            return False
        
        try:
            with open(path_to_check, 'r', encoding='utf-8') as f:
                knowledge = json.load(f)
            
            if not isinstance(knowledge, list):
                print("错误：知识库必须是一个数组")
                return False
            
            required_fields = ["id", "title", "content", "code", "tags", "category", "difficulty"]
            seen_ids = set()
            invalid_entries = []
            
            for i, entry in enumerate(knowledge):
                # 检查必需字段
                missing_fields = [field for field in required_fields if field not in entry]
                if missing_fields:
                    invalid_entries.append((i, f"缺少必需字段: {', '.join(missing_fields)}"))
                    continue
                
                # 检查ID唯一性
                if entry["id"] in seen_ids:
                    invalid_entries.append((i, f"ID重复: {entry['id']}"))
                seen_ids.add(entry["id"])
                
                # 检查字段类型
                if not isinstance(entry["tags"], list):
                    invalid_entries.append((i, "tags必须是数组类型"))
                
            if invalid_entries:
                print(f"发现{len(invalid_entries)}个无效条目:")
                for idx, error in invalid_entries:
                    print(f"  条目{idx}: {error}")
                return False
            
            print(f"知识库验证成功，共{len(knowledge)}条有效记录")
            return True
            
        except json.JSONDecodeError as e:
            print(f"JSON格式错误: {e}")
            return False
        except Exception as e:
            print(f"验证过程中发生错误: {e}")
            return False

def main():
    print("=== Rust知识库扩充工具 ===")
    print("此工具将帮助您自动扩充Rust知识库，无需手动输入\n")
    
    # 创建扩充器实例
    expander = RustKnowledgeBaseExpander()
    
    # 加载现有知识库
    expander.load_existing_knowledge()
    
    # 设置目标数量
    target_count = 500  # 默认目标500条
    
    # 生成新条目
    expander.generate_knowledge_base(target_count)
    
    # 合并并保存
    expander.combine_and_save()
    
    # 验证生成的知识库
    expander.validate_knowledge_base()
    
    print("\n=== 知识库扩充完成 ===")
    print("您可以在config.py中修改KNOWLEDGE_BASE_PATH指向新的知识库文件")
    print("或者将生成的文件复制到rust_knowledge_base/rust_docs_sample.json替换原有文件")

if __name__ == "__main__":
    main()